import numpy as np
import datetime
import pandas as pd


def polyfit(x, y, degree):
    results = {}
    coeffs = np.polyfit(x, y, degree)
    results["polynomial"] = coeffs.tolist()

    # r-squared
    p = np.poly1d(coeffs)
    Q = np.polyder(coeffs)  # 求导数
    xs = np.roots(Q)  # 求多项式函数根
    ys = np.polyval(coeffs, xs)
    # fit values, and mean

    yhat = p(x)  # or [p(z) for z in x]
    ybar = np.sum(y) / len(y)  # or sum(y)/len(y)
    ssreg = np.sum((yhat - ybar) ** 2)  # or sum([ (yihat - ybar)**2 for yihat in yhat])
    sstot = np.sum((y - ybar) ** 2)  # or sum([ (yi - ybar)**2 for yi in y])

    newX = []
    for i in range(len(x) + 5):
        newX.append(float(i))
    newYhat = p(newX)
    if sstot != 0:
        results["determination"] = ssreg / sstot  # 准确率
    else:
        results["determination"] = 0
    results["Y"] = newYhat
    results["Q"] = Q
    results["xs"] = xs
    results["ys"] = ys
    return results


def polyfitlize(df, key, gap=5):
    """
    对数据进行拟合
    """
    l = []
    XYl = []
    YDic = {}
    pre_determination = []
    start = 0
    if gap > 40:
        ranges = gap * 2
    elif gap > 10:
        ranges = gap * 3
    else:
        ranges = gap * 4
    starttime = datetime.datetime.now()
    print("{} polyfitlize start".format(key), end=" ")
    total = df.shape[0] - ranges
    prescale = 0
    for _ in range(total):
        scale = int((_ / total) * 100)
        if scale % 5 == 0 and prescale != scale:
            print("*", end="")
        prescale = scale

        y = df.iloc[start + _ : ranges + _][key]
        y.dropna(axis=0, inplace=True)
        while np.nan in y:
            y.remove(np.nan)

        dic = {}
        dic["date"] = df.iloc[_]["date"]
        XYdic = {}
        x = []
        a = 0
        for i in y:
            a += 1
            x.append(float(a))

        if len(x) >= ranges and x != []:
            results = polyfit(x, y, 3)
            dic["determination"] = results["determination"]
            a = 0
            coeffsL = results["polynomial"]
            for i in coeffsL:
                dic["polynomial_X_{}".format(a)] = coeffsL[a]
                a += 1

            a = 0
            Ql = results["Q"]
            for i in Ql:
                dic["Q_{}".format(a)] = Ql[a]
                a += 1

            if (
                0.87 < results["determination"] < 1
            ):  # TODO 添加对已存在点的 determination比较,经可能平滑曲线
                a = 0
                yL = results["Y"]
                for i in yL:
                    YDic[df.iloc[_ + a].date] = [round(i, 5), df.iloc[_].date]
                    a += 1

            if 0.90 < results["determination"] < 0.999:  # TODO 合并到上一个循环中
                a = 0
                xsL = results["xs"].tolist()
                ysL = results["ys"].tolist()
                for xs in xsL:
                    if "j" in str(xs):
                        temp = str(xs)
                        try:
                            if "+" in temp or temp.startswith("(-"):
                                _xs = temp.split("+")[0]
                                if "-" in _xs:
                                    _xs = _xs.split("-")[0]
                                    _xs = -round(float(_xs.replace("(", "0")))
                                else:
                                    _xs = round(float(_xs.replace("(", "0")))

                            elif "-" in temp:
                                _xs = (
                                    round(float((temp.split("-")[0]).replace("(", "0")))
                                    + _
                                )
                        except Exception as e:
                            print("error --> {} {}".format(temp, xs))
                            print(e)
                    else:
                        _xs = float(xs) + _

                    if 0 < _xs < df.shape[0] and _xs:
                        gap_days = datetime.datetime.strptime(
                            df.iloc[_].date, "%Y-%m-%d %H:%M:%S"
                        ) - datetime.datetime.strptime(
                            df.iloc[int(_xs)].date, "%Y-%m-%d %H:%M:%S"
                        )
                        if abs(gap_days.days) * gap <= ranges:
                            XYdic["date"] = df.iloc[_].date
                            XYdic["xss_{}".format(a)] = df.iloc[int(_xs)].date
                            ys = ysL[xsL.index(xs)]
                            if "j" in str(ys):
                                temp = str(ys)
                                try:
                                    if "+" in temp or temp.startswith("-"):
                                        yst = temp.split("+")[0]
                                        if "-" in yst:
                                            yst = yst.split("-")[0]
                                            ys = -float(yst.replace("(", "0"))
                                        else:
                                            ys = float(yst.replace("(", "0"))
                                    elif "-" in temp:
                                        ys = float(
                                            (temp.split("-")[0]).replace("(", "0")
                                        )
                                except Exception as e:
                                    print("error --> {} {}".format(temp, ys))
                                    print(e)
                            else:
                                ys = float(ys)
                            XYdic["ys_{}".format(a)] = round(ys, 4)

                    a += 1
                XYl.append(XYdic)
        l.append(dic)  # TODO 修改信息加载方式,可以通过concat将数据进行合并

    if l != []:
        df2 = pd.DataFrame(l)
        df = pd.merge(df, df2, on="date", how="left")

    if YDic != {}:
        _dfY = pd.DataFrame(YDic, index=["Y_prodict", "prodict_date"]).T.reset_index()
        _dfY.rename(columns={"index": "date"}, inplace=True)
        df = pd.merge(df, _dfY, on="date", how="left")

    if XYl != []:
        _dfX = pd.DataFrame(XYl)
        df = pd.merge(df, _dfX, on="date", how="left")

    print(" polyfitlize => {}s".format((datetime.datetime.now() - starttime).seconds))
    return df